import torchvision
import torch
from torch import nn,optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision import transforms
import torch.nn.functional as F
import numpy as np
import os

Batch_size = 128
#### 准备数据集
def get_dataloader(train=True):
    transform_fn = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.1307,),std=(0.3081,))
    ])
    mnist = MNIST(root='./data/',train=train,transform=transform_fn)
    data_loader = DataLoader(dataset=mnist,batch_size=Batch_size,shuffle=True,drop_last=True)
    return data_loader


#### 构建模型
class Mnist_Module(nn.Module):
    def __init__(self):
        super(Mnist_Module,self).__init__()
        self.fc1 = nn.Linear(1*28*28,28)
        self.fc2 = nn.Linear(28,10)

    def forward(self,input):
        ### 进行全连接之前，修改形状变为二维数据
        x = input.view([Batch_size,1*28*28])

        x = self.fc1(x)
        x = F.relu(x)
        out = self.fc2(x)
        return F.log_softmax(out,dim=-1)

#### 实例化模型
model = Mnist_Module()
optimizer = optim.Adam(model.parameters(), lr=0.001)

if os.path.exists('./model/model.pkl'):
    model.load_state_dict(torch.load('./model/model.pkl'))
    optimizer.load_state_dict(torch.load(('./model/optimizer.pkl')))



def train(epoch):
    data_loader = get_dataloader()
    for idx,(input,target) in enumerate(data_loader):
        output = model(input)
        loss = F.nll_loss(output,target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if idx % 100==0:
            print('epoch:{},idx:{},loss:{}'.format(epoch,idx,loss))
            torch.save(model.state_dict(),'./model/model.pkl')
            torch.save(optimizer.state_dict(),'./model/optimizer.pkl')


def test():
    loss_list = []
    acc_list = []
    model.eval()
    test_data_loader = get_dataloader(train=False)
    for idx,(input,target) in enumerate(test_data_loader):
        with torch.no_grad():
            output = model(input)
            test_loss = F.nll_loss(output,target)
            loss_list.append(test_loss)
            ### 计算准确率
            pred = output.max(dim=-1)[-1]
            cur_acc = pred.eq(target).float().mean()
            acc_list.append(cur_acc)

    print('平均准确率：{},平均损失：{}'.format(np.mean(acc_list),np.mean(loss_list)))


if __name__ == '__main__':
    for i in range(3):
        # train(i)
        test()
